# 手动进行问答，不使用langchain库(限制太多)
from abc import ABC

#from config.config_lab import is_llm_ie, is_chat_kg
from question.kb_service.milvus_service import MilvusService
from question.llm_service.chatglmService import ChatGLMService
from question.llm_service.qwenService import QwenService
from question.kb_service.faiss_service import KGService
from question.chat_chain.config import CFG
from langchain import LLMChain, PromptTemplate

#from utils import Neo4jUtil, UIEUtil, LLMIEUtil


class ChatByKB(ABC):
    def __init__(self):
        self.config = CFG()
        if "chatglm" in self.config.llm_model_name:
            self.llm_service = ChatGLMService()
            self.llm_service.load_model(model_name_or_path=self.config.llm_model_name
                                        , isqiamtize4=False, device=self.config.device)
        elif "Qwen" in self.config.llm_model_name:
            self.llm_service = QwenService()
            self.llm_service.load_model(model_name_or_path=self.config.llm_model_name
                                        , isqiamtize4=False, device=self.config.device)

        if self.config.knowledge_base_type == "faiss":
            self.source_service = KGService(knowledge_base_name=self.config.knowledge_base_name)
        elif self.config.knowledge_base_type == "Milvus":
            self.source_service = MilvusService(knowledge_base_name=self.config.knowledge_base_name)
        self.source_service.load_vector_store()
        # 初始化知识图谱服务
        # 延迟导入 utils 中的类（解决循环导入）
        from utils.Neo4jUtil import Neo4jUtil
        from utils.UIEUtil import UIEUtil

        self.neo4j_util = Neo4jUtil()

        # 初始化实体抽取工具
        self.uie_util = UIEUtil()
        self.uie_util.create_ie_instance()
        #self.llm_ie = LLMIEUtil()
        print("QA初始化完成!")

    def extract_entities(self, query):
        """从问题中抽取实体-关系对（实体+关联的关系）"""
        predict = self.uie_util.predict(query)
        entity_relation_pairs = []  # 存储(实体, 关系)元组
        for p_item in predict:
            for item_key in p_item.keys():
                item_data = p_item[item_key]
                for item in item_data:
                    entity = item['text']
                    # 提取该实体关联的关系
                    if 'relations' in item:
                        relations = item['relations'].keys()  # 关系类型（如['包含']）
                        for rel in relations:
                            entity_relation_pairs.append((entity, rel))
                    else:
                        # 若无明确关系，可留空或按原逻辑处理
                        entity_relation_pairs.append((entity, ""))
        # 去重
        entity_relation_pairs = list(set(entity_relation_pairs))
        return entity_relation_pairs

    def get_knowledge_graph_info(self, entity_relation_pairs, level=2):
        """从知识图谱中获取实体+关系相关的信息"""
        graph_info = ""
        for entity, relation in entity_relation_pairs:
            # 若存在关系，按“实体+关系”查询；否则按实体查询（兼容旧逻辑）
            if relation:
                # 调用search_all，传入entity1和relation，查询指定关系的结果
                result = self.neo4j_util.search_all(
                    entity1=entity,
                    relation=relation,
                    numberOfLevels=level
                )
            else:
                # 无关系时，查询实体所有关联信息（保持原逻辑）
                result = self.neo4j_util.search_all(
                    entity1=entity,
                    numberOfLevels=level
                )

            if not result:
                continue

            graph_info += f"\n与实体「{entity}」（关系：{relation}）相关的知识图谱信息：\n"
            nodes, edges = self.search_res_path_to_graph(result)

            # 格式化节点信息（同原逻辑）
            node_info = {node["id"]: node["cluster"] for node in nodes}
            for edge in edges:
                source = edge["source"]
                target = edge["target"]
                rel = edge["relation"]
                graph_info += f"{source}（{node_info[source]}）{rel} {target}（{node_info[target]}）\n"

        return graph_info

    def search_res_path_to_graph(self, searchResult):
        """将图谱查询结果转换为节点和边的格式（兼容path和n1/rel/n2两种结构）"""
        if searchResult is None or len(searchResult) == 0:
            return [], []

        nodesRes = list()
        edgesRes = list()
        nodeSet = set()
        edgeSet = set()

        for item in searchResult:
            # 处理包含'path'字段的结果（多级路径查询）
            if 'path' in item:
                path = item['path']
                nodes = [node['name'] for node in path.nodes]
                rels = [list(rel.types())[0] for rel in path.relationships]
                node_types = [str(node.labels).strip(":") for node in path.nodes]
                # 提取节点
                for i in range(len(nodes)):
                    nodeSet.add((nodes[i], node_types[i]))
                # 提取关系
                for i in range(len(rels)):
                    edgeSet.add((nodes[i], rels[i], nodes[i + 1]))

            # 处理包含'n1'/'rel'/'n2'字段的结果（实体+关系查询）
            elif 'n1' in item and 'rel' in item and 'n2' in item:
                n1 = item['n1']
                rel = item['rel']
                n2 = item['n2']
                # 节点信息
                node1_name = n1['name']
                node1_type = str(n1.labels).strip(":")
                node2_name = n2['name']
                node2_type = str(n2.labels).strip(":")
                # 关系信息
                relation = list(rel.types())[0]
                # 添加节点和关系
                nodeSet.add((node1_name, node1_type))
                nodeSet.add((node2_name, node2_type))
                edgeSet.add((node1_name, relation, node2_name))

        # 转换为节点列表
        for node in nodeSet:
            nodesRes.append({
                "id": node[0],
                "label": node[0],
                "cluster": node[1]
            })
        # 转换为边列表
        for edge in edgeSet:
            edgesRes.append({
                "source": edge[0],
                "target": edge[2],
                'relation': edge[1]
            })

        return nodesRes, edgesRes
    def get_knowledge_based_answer(self, query,is_chat_kg=False,
                                   temperature=0.7,
                                   history_len=0,
                                   top_p=0.9,
                                   top_k=3,
                                   chat_history=[],
                                   score_threshold=0.5):
        docs = self.source_service.do_search(query=query, top_k=top_k, score_threshold=score_threshold)
        print("docs:",docs)
        source_documents = ""
        for inum, doc in enumerate(docs):
            filename = doc[0].metadata.get("source")
            text = f"""出处 [{inum + 1}] [{filename}], 内容: \n\n{doc[0].page_content}\n\n"""
            source_documents = source_documents + text

        graph_info={}
        if is_chat_kg:
            # 2. 从知识图谱获取相关信息
            entity_relation_pairs = self.extract_entities(query)
            print(f"从问题中抽取的实体-关系对: {entity_relation_pairs}")
            # 传入实体-关系对查询图谱
            graph_info = self.get_knowledge_graph_info(entity_relation_pairs)
            print(f"知识图谱信息: {graph_info}")

            prompt_template = {
                "default":
                    '<指令>根据已知信息，简洁和专业的来回答问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题”，'
                    '不允许在答案中添加编造成分，答案请使用中文。 </指令>\n'
                    '<已知信息>{context}</已知信息>\n'
                    '<相关知识图谱信息>{graph_context}</相关知识图谱信息>\n'
                ,
                "empty":  # 搜不到知识库的时候使用
                    '请你回答我的问题:\n'
                    '{question}\n\n',
            }
            prompt_add = '<问题>{question}</问题>\n'
        else:
            prompt_template = {
                "default":
                    '<指令>根据已知信息，简洁和专业的来回答问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题”，'
                    '不允许在答案中添加编造成分，答案请使用中文。 </指令>\n'
                    '<已知信息>{context}</已知信息>\n'
                ,
                "empty":  # 搜不到知识库的时候使用
                    '请你回答我的问题:\n'
                    '{question}\n\n',
            }
            prompt_add = '<问题>{question}</问题>\n'
        self.llm_service.temperature = temperature
        self.llm_service.top_p = top_p
        self.llm_service.history = chat_history[-history_len:] if history_len > 0 else []
        if len(source_documents) != 0:
            # 查询到知识则使用知识问答
            if is_chat_kg:
                prompt_template['default'] = prompt_template['default'].format(context=source_documents,graph_context=graph_info)
            else:
                prompt_template['default'] = prompt_template['default'].format(context=source_documents)
            prompt_template['default'] = prompt_template['default'] + prompt_add
            prompt = PromptTemplate(template=prompt_template['default'],
                                    input_variables=["question"])
            print("prompt",prompt)
            chain = LLMChain(
                llm=self.llm_service,
                prompt=prompt
            )
            result = chain({"question": query})
            return result['text']
        else:
            # 没有知识则使用大语言模型自身能力
            prompt = PromptTemplate(template=prompt_template['empty'],
                                    input_variables=["question"])
            chain = LLMChain(
                llm=self.llm_service,
                prompt=prompt
            )
            result = chain({"question": query})
            return result['text']
